package com.yolov8.steam;

import com.yolov8.model.FrameData;
import lombok.extern.slf4j.Slf4j;
import org.bytedeco.javacv.FFmpegFrameGrabber;
import org.bytedeco.javacv.Frame;
import org.bytedeco.javacv.OpenCVFrameConverter;
import org.bytedeco.opencv.opencv_core.Mat;

import java.util.concurrent.BlockingQueue;

@Slf4j
public class VideoPuller implements Runnable {

    private final String url;
    private final BlockingQueue<FrameData> pullQueue;
    private final int maxRetry;
    private volatile boolean running = true;

    public VideoPuller(String url, BlockingQueue<FrameData> pullQueue, int maxRetry) {
        this.url = url;
        this.pullQueue = pullQueue;
        this.maxRetry = maxRetry;
    }

    @Override
    public void run() {
        int retryCount = 0;
        while (running) {
            try (FFmpegFrameGrabber grabber = new FFmpegFrameGrabber(url)) {
                grabber.start();
                log.info("开始拉流: " + url);

                Frame frame;
                OpenCVFrameConverter.ToMat converter = new OpenCVFrameConverter.ToMat();

                int num = 0;
                while (running && (frame = grabber.grabImage()) != null) {
                    Mat mat = converter.convert(frame);
                    if (mat != null) {
                        num++;
                        pullQueue.put(new FrameData(mat.clone(), System.currentTimeMillis()));
                        log.info("第 "+num+" 次拉到帧: " + mat.rows() + " x " + mat.cols());
                        log.info("拉流队列大小"+ pullQueue.size());
//                        log.info("第 "+num+" 次拉到帧: ");
                    }
                }

                grabber.stop();
                break;
            } catch (Exception e) {
                retryCount++;

                log.info("拉流失败: " + e.getMessage());
                if (retryCount >= maxRetry) {
                    log.info("重试次数过多，终止任务");
                    stop();
                } else {
                    try {
                        Thread.sleep(2000); // 等待后重试
                    } catch (InterruptedException ignored) {}
                }
            }
        }
    }

    public void stop() {
        this.running = false;
    }
}
